package com.hxl.tech.gateway.core.filter;

import cn.hutool.core.collection.CollUtil;
import com.hxl.tech.gateway.common.constant.AppConstant;
import com.hxl.tech.gateway.common.dto.LimitRuleDTO;
import com.hxl.tech.gateway.core.util.CacheManager;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

@Slf4j
@Component
@RequiredArgsConstructor
public class LimitRateFilter implements GlobalFilter, Ordered {

    private final Map<Long, Map<String, AtomicInteger>> rateLimitMap = new ConcurrentHashMap<>();

    private final CacheManager cacheManager;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        return Mono.subscriberContext().flatMap(ctx -> {
            String appId = ctx.get(AppConstant.APP_ID);
            String apiName = ctx.get(AppConstant.API_NAME);
            String merchantCode = ctx.get(AppConstant.MERCHANT_CODE);

            List<LimitRuleDTO> limitRules = cacheManager.getRateLimit(appId, merchantCode, apiName);
            if (CollUtil.isEmpty(limitRules)) {
                return chain.filter(exchange);
            }

            Map<String, List<LimitRuleDTO>> map = limitRules.stream().collect(Collectors.groupingBy(LimitRuleDTO::getLimitTimeUnit));

            // 针对每秒得处理
            List<LimitRuleDTO> limitRuleDTOSBySecond = map.get("秒");
            if (CollUtil.isNotEmpty(limitRuleDTOSBySecond)) {
                String path = appId + "_" + merchantCode + "_" + apiName + "_s";
                long nowSecond = System.currentTimeMillis() / 1000;
                Map<String, AtomicInteger> apiCount = rateLimitMap.computeIfAbsent(nowSecond, k->new HashMap<String, AtomicInteger>(){{
                    put(path, new AtomicInteger(1));
                }});

                AtomicInteger count = apiCount.computeIfAbsent(path, k->new AtomicInteger(1));
                int countValue = count.get();
                log.debug("请求主体:[{}], 每秒钟请求次数:[{}]", path, countValue);
                if(countValue > limitRuleDTOSBySecond.get(0).getLimitCount()) {
                    log.info("网关限流策略:[{}]，超出每秒请求次数:[{}]次", path, limitRuleDTOSBySecond.get(0).getLimitCount());
                    exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                    return exchange.getResponse().setComplete();
                }
                count.addAndGet(1);
                return chain.filter(exchange).then(Mono.fromRunnable(()->{
                    rateLimitMap.keySet().stream().filter(time->time<nowSecond).forEach(rateLimitMap::remove);
                }));
            }

            // 针对每分得处理
            List<LimitRuleDTO> limitRuleDTOSByMinute = map.get("分");
            if (CollUtil.isNotEmpty(limitRuleDTOSByMinute)) {
                String path = appId + "_" + merchantCode + "_" + apiName + "_m";
                long nowMinute = System.currentTimeMillis() / 1000 / 60;
                Map<String, AtomicInteger> apiCount = rateLimitMap.computeIfAbsent(nowMinute, k->new HashMap<String, AtomicInteger>(){{
                    put(path, new AtomicInteger(1));
                }});

                AtomicInteger count = apiCount.computeIfAbsent(path, k->new AtomicInteger(1));
                int countValue = count.get();
                log.debug("请求主体:[{}], 每分钟请求次数:[{}]", path, countValue);
                if(countValue > limitRuleDTOSByMinute.get(0).getLimitCount()) {
                    log.info("网关限流策略:[{}]，超出每分请求次数:[{}]次", path, limitRuleDTOSByMinute.get(0).getLimitCount());
                    exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                    return exchange.getResponse().setComplete();
                }
                count.addAndGet(1);
                return chain.filter(exchange).then(Mono.fromRunnable(()->{
                    rateLimitMap.keySet().stream().filter(time->time<nowMinute).forEach(rateLimitMap::remove);
                }));
            }
            return chain.filter(exchange);
        });
    }

    @Override
    public int getOrder() {
        return 5002;
    }
}
